"""
运行环境：
pip install flask pillow numpy onnxruntime
"""


from flask import Flask, render_template, request, jsonify
import os
from PIL import Image
import numpy as np
import onnxruntime as ort
import base64
from io import BytesIO
import json

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB限制

# 全局变量存储当前模型
current_model = {
    'session': None,
    'input_name': None,
    'output_name': None,
    'class_names': [],
    'model_path': '',
    'label_path': '',
    'input_shape': None
}

def load_model(model_path, label_path):
    """加载ONNX模型和标签"""
    try:
        # 加载模型
        sess = ort.InferenceSession(model_path)
        input_name = sess.get_inputs()[0].name
        output_name = sess.get_outputs()[0].name
        input_shape = sess.get_inputs()[0].shape
        
        # 加载标签
        with open(label_path, 'r', encoding='utf-8') as f:
            class_names = [line.strip() for line in f.readlines()]
        
        # 更新全局状态
        current_model.update({
            'session': sess,
            'input_name': input_name,
            'output_name': output_name,
            'class_names': class_names,
            'model_path': model_path,
            'label_path': label_path,
            'input_shape': input_shape
        })
        
        return True, "模型加载成功"
    except Exception as e:
        current_model['session'] = None
        return False, f"加载失败: {str(e)}"

# def preprocess_image(image):
#     """图片预处理"""
#     img = Image.open(BytesIO(image))
#     img = img.resize((224, 224))  # 根据模型输入调整
#     img = np.array(img).astype(np.float32)
#     img = img / 255.0  # 归一化
#     return np.expand_dims(img, axis=0)  # 添加batch维度

def preprocess_image(image_bytes):
        """预处理图片以适应模型输入"""
        try:
            img = Image.open(BytesIO(image_bytes))
            
            # 根据模型输入形状调整大小
            input_shape = current_model['input_shape']
            if len(input_shape) == 4:  # NCHW格式
                input_height, input_width = input_shape[2], input_shape[3]
            else:  # 假设是NHWC格式
                input_height, input_width = input_shape[1], input_shape[2]
            
            img = img.resize((input_width, input_height))
            
            # 转换为numpy数组并归一化
            img_array = np.array(img).astype(np.float32)
            img_array = img_array / 255.0  # 归一化
            
            # 添加批次维度
            if len(input_shape) == 4:
                if input_shape[1] == 3:  # 假设是RGB通道在前
                    img_array = img_array.transpose((2, 0, 1))  # HWC to CHW
                img_array = np.expand_dims(img_array, axis=0)  # 添加批次维度
            
            return img_array
        except Exception as e:
            print(f"图片预处理失败: {str(e)}")  # 打印错误信息
            return None

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # 处理模型加载请求
        if 'model_path' in request.form:
            model_path = request.form['model_path']
            label_path = request.form['label_path']
            
            success, message = load_model(model_path, label_path)
            return jsonify({"success": success, "message": message})
        
        # 处理图片识别请求
        elif 'file' in request.files:
            if not current_model['session']:
                return jsonify({"error": "请先加载模型文件"}), 400
                
            file = request.files['file'].read()
            if not file:
                return jsonify({"error": "空文件"}), 400

            try:
                # 预处理和推理
                input_tensor = preprocess_image(file)
                results = current_model['session'].run(
                    [current_model['output_name']], 
                    {current_model['input_name']: input_tensor}
                )[0]
                pred_idx = np.argmax(results)
                
                # 转换为base64预览图
                img = Image.open(BytesIO(file)).resize((300, 300))
                buffered = BytesIO()
                img.save(buffered, format="JPEG")
                img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
                
                return jsonify({
                    "preview": f"data:image/jpeg;base64,{img_str}",
                    "prediction": current_model['class_names'][pred_idx],
                    "confidence": float(results[0][pred_idx]),
                    "top3": [
                        {"class": current_model['class_names'][i], "score": float(score)} 
                        for i, score in enumerate(results[0])
                    ][:3]
                })
                
            except Exception as e:
                return jsonify({"error": f"识别错误: {str(e)}"}), 500
    
    return render_template('index.html')

if __name__ == '__main__':
    # 尝试加载上次使用的模型（如果存在）
    if os.path.exists('last_model.json'):
        with open('last_model.json') as f:
            config = json.load(f)
            if os.path.exists(config['model_path']) and os.path.exists(config['label_path']):
                load_model(config['model_path'], config['label_path'])
    
    app.run(debug=True)